# -*- coding: utf-8 -*-
import cv2
import numpy as np
import mediapipe as mp

mp_pose = mp.solutions.pose
LMS = mp_pose.PoseLandmark

# 画骨架的连线对
BONES = [
    (LMS.LEFT_SHOULDER, LMS.RIGHT_SHOULDER),
    (LMS.LEFT_SHOULDER, LMS.LEFT_ELBOW), (LMS.LEFT_ELBOW, LMS.LEFT_WRIST),
    (LMS.RIGHT_SHOULDER, LMS.RIGHT_ELBOW), (LMS.RIGHT_ELBOW, LMS.RIGHT_WRIST),
    (LMS.LEFT_HIP, LMS.RIGHT_HIP),
    (LMS.LEFT_SHOULDER, LMS.LEFT_HIP), (LMS.RIGHT_SHOULDER, LMS.RIGHT_HIP),
    (LMS.LEFT_HIP, LMS.LEFT_KNEE), (LMS.LEFT_KNEE, LMS.LEFT_ANKLE),
    (LMS.RIGHT_HIP, LMS.RIGHT_KNEE), (LMS.RIGHT_KNEE, LMS.RIGHT_ANKLE),
]

def draw_landmarks(frame_bgr, landmarks, visibility_thr=0.3, point_color=(0, 255, 0)):
    h, w = frame_bgr.shape[:2]
    # 画线
    for a, b in BONES:
        la = landmarks[a.value]
        lb = landmarks[b.value]
        if getattr(la, "visibility", 0.0) < visibility_thr or getattr(lb, "visibility", 0.0) < visibility_thr:
            color = (150, 150, 150)
        else:
            color = (200, 200, 200)
        pa = (int(la.x * w), int(la.y * h))
        pb = (int(lb.x * w), int(lb.y * h))
        cv2.line(frame_bgr, pa, pb, color, 2)

    # 画点
    for i, l in enumerate(landmarks):
        color = (0, 255, 0) if getattr(l, "visibility", 0.0) >= visibility_thr else (128, 128, 128)
        cv2.circle(frame_bgr, (int(l.x * w), int(l.y * h)), 4, color, -1)
        